import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, LabelEncoder
from sklearn.neighbors import KNeighborsClassifier
from sklearn.svm import SVC
from sklearn.metrics import classification_report, accuracy_score
from joblib import dump

class Classification(object):
    """股票分类与预测模型工具类"""
    
    def __init__(self):
        """初始化并加载数据"""
        self.df = pd.read_csv('daily.csv')
    
    def get_conditions(self):
        """
        数据预处理与特征工程
        返回:
            X_train: 训练集特征
            X_test: 测试集特征
            y_train: 训练集标签
            y_test: 测试集标签
            le: 标签编码器对象
        """
        df = self.df.copy()
        
        # 计算收益和风险比例
        df['max_ratio'] = df['max_close'] / df['the_close']  # 收益能力指标
        df['min_ratio'] = df['min_close'] / df['the_close']  # 风险控制指标
        
        # 自动确定阈值 - 使用分位数确定高低收益/风险的界限
        high_return_threshold = df['max_ratio'].quantile(0.6)  # 前40%定义为高收益
        high_risk_threshold = df['min_ratio'].quantile(0.4)    # 后40%定义为高风险
        
        # 定义分类条件
        conditions = [
            (df['max_ratio'] >= high_return_threshold) & (df['min_ratio'] <= high_risk_threshold),
            (df['max_ratio'] >= high_return_threshold) & (df['min_ratio'] > high_risk_threshold),
            (df['max_ratio'] < high_return_threshold) & (df['min_ratio'] > high_risk_threshold),
            (df['max_ratio'] < high_return_threshold) & (df['min_ratio'] <= high_risk_threshold)
        ]
        
        # 定义分类标签
        labels = ['高收益高风险', '高收益低风险', '低收益低风险', '低收益高风险']
        df['category'] = np.select(conditions, labels, default='未知')
        
        # 选择特征列
        features = df[['eps', 'total_revenue_ps', 'undist_profit_ps', 
                      'gross_margin', 'fcff', 'fcfe', 'tangible_asset',
                      'bps', 'grossprofit_margin', 'npta']]
        
        # 标签编码
        le = LabelEncoder()
        df['category_encoded'] = le.fit_transform(df['category'])
        
        # 数据标准化
        scaler = StandardScaler()
        X = scaler.fit_transform(features)
        y = df['category_encoded']
        
        # 划分训练集和测试集
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=0.3, random_state=24
        )
        
        return X_train, X_test, y_train, y_test, le
    
    def knn_utils(self, X_train, X_test, y_train, y_test, le):
        """
        K近邻分类模型训练与评估
        参数:
            X_train: 训练集特征
            X_test: 测试集特征
            y_train: 训练集标签
            y_test: 测试集标签
            le: 标签编码器对象
        """
        # 初始化KNN模型
        knn = KNeighborsClassifier(n_neighbors=3)
        knn.fit(X_train, y_train)
        
        # 预测与评估
        y_pred = knn.predict(X_test)
        
        print("KNN分类结果报告:")
        print(classification_report(
            y_test, y_pred, 
            target_names=le.classes_
        ))
        
        # 保存模型
        dump(knn, 'knn_model.joblib')
        print("KNN模型已保存为knn_model.joblib")
    
    def svc_utils(self, X_train, X_test, y_train, y_test):
        """
        支持向量机模型训练与评估
        参数:
            X_train: 训练集特征
            X_test: 测试集特征
            y_train: 训练集标签
            y_test: 测试集标签
        """
        # 初始化SVM模型
        svc = SVC()
        svc.fit(X_train, y_train)
        
        # 预测与评估
        predict = svc.predict(X_test)
        
        print("\nSVM分类结果:")
        print("准确率: %.4f" % accuracy_score(predict, y_test))
        print("分类报告:\n", classification_report(
            y_test, predict
        ))
        
        # 保存模型
        dump(svc, 'svc_model.joblib')
        print("SVM模型已保存为svc_model.joblib")

if __name__ == '__main__':
    # 主程序入口
    cl = Classification()
    
    # 数据预处理
    X_train, X_test, y_train, y_test, le = cl.get_conditions()
    
    # 运行KNN模型
    print("\n" + "="*50)
    print("正在运行KNN分类模型...")
    cl.knn_utils(X_train, X_test, y_train, y_test, le)
    
    # 运行SVM模型
    print("\n" + "="*50)
    print("正在运行SVM分类模型...")
    cl.svc_utils(X_train, X_test, y_train, y_test)